import subprocess
from typing import Any, Dict, Iterator, Optional

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults


class LogstashOperator(BaseOperator):

    @apply_defaults
    def __init__(
            self,
            *,
            content: Optional[str] = None,
            logstash_binary: Optional[str] = None,
            **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._content = content
        self._logstash_binary = logstash_binary
        self._proc = None

    def execute(self, context: Dict[str, Any]) -> None:
        self.submit(self._content)

    def on_kill(self) -> None:
        """Kill Spark submit command"""
        self.log.debug("Kill Command is being called")

    def _build_logstash_command(self, content):
        fn = '/opt/airflow/logstash.conf'
        with open(fn, 'wt', encoding='utf-8') as fin:
            fin.write(content)
        cmd = []
        if self._logstash_binary:
            cmd.append(self._logstash_binary)
        else:
            cmd.append('/opt/logstash/bin/logstash')
        cmd.append('-f')
        cmd.append(fn)
        return cmd

    def submit(self, content: str = "", **kwargs: Any) -> None:
        """
        Remote Popen to execute the flink-submit job

        :param application: Submitted application, jar or py file
        :type application: str
        :param kwargs: extra arguments to Popen (see subprocess.Popen)
        """
        logstash_cmd = self._build_logstash_command(content)
        self._proc = subprocess.Popen(
            logstash_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            bufsize=-1,
            universal_newlines=True
        )

        self._process_logstash_log(iter(self._proc.stdout))  # type: ignore
        return_code = self._proc.wait()
        self.log.info('logstash exit {}'.format(return_code))

    def _process_logstash_log(self, itr: Iterator[Any]) -> None:
        """
        Processes the log files and extracts useful information out of it.
        :param itr: An iterator which iterates over the input of the subprocess
        """
        # Consume the iterator
        for line in itr:
            line = line.strip()
            self.log.info(line)
